import os
import pinecone
from pinecone import Pinecone, ServerlessSpec
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score
import logging
from tqdm import tqdm
import joblib  # 导入 joblib 库

# 设置日志格式，包含日期和时间
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 使用 Pinecone 实例化
api_key = "8426bf47-0774-4cb3-b4b2-7347f267a02e"  # 将此替换为你的实际 API key
environment = "us-east-1"  # 替换为你的实际环境
pinecone_instance = Pinecone(api_key=api_key)

# 创建 Pinecone 索引
index_name = "mnist-index"
if index_name not in pinecone_instance.list_indexes().names():
    pinecone_instance.create_index(
        name=index_name,
        dimension=64,
        metric='euclidean',
        spec=ServerlessSpec(cloud='aws', region=environment)
    )

index = pinecone_instance.Index(index_name)

# 加载 MNIST 数据集
digits = load_digits()
X = digits.data
y = digits.target

# 数据标准化
scaler = StandardScaler()
X = scaler.fit_transform(X)

# 将数据集拆分为训练集（80%）和测试集（20%）
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 上传训练数据到 Pinecone，使用进度条显示上传进度
batch_size = 1000  # 每次最多上传1000条数据
data_to_upload = []
for i, (vector, label) in enumerate(zip(X_train, y_train)):
    data_to_upload.append((str(i), vector.tolist(), {"label": int(label)}))
    
    # 当达到批次大小时，上传数据
    if (i + 1) % batch_size == 0 or i == len(X_train) - 1:
        index.upsert(vectors=data_to_upload)
        logging.info(f"已上传 {i + 1} 条数据")
        data_to_upload = []  # 清空缓冲区，准备下一批数据

logging.info(f"成功创建索引，并上传了{len(X_train)}条数据")

# 测试阶段，使用KNN（k=11）
k = 11
knn = KNeighborsClassifier(n_neighbors=k)
knn.fit(X_train, y_train)

# 保存最佳KNN模型到二进制文件
model_filename = 'best_knn_model.pkl'
joblib.dump(knn, model_filename)
logging.info(f"最佳KNN模型已保存到 {model_filename}")

# 测试并记录准确率，使用进度条显示测试进度
y_pred = []
for i, test_vector in tqdm(enumerate(X_test), total=len(X_test), desc="Testing k=11"):
    test_vector = test_vector.reshape(1, -1)
    pred = knn.predict(test_vector)
    y_pred.append(pred[0])

accuracy = accuracy_score(y_test, y_pred)
logging.info(f"当k={k}时，使用Pinecone的准确率: {accuracy * 100:.2f}%")

